import os
from fasttext import FTModel
from create_alarms_data_002 import ALL_ORIGIN_INFOS
import torch
from dataset_004 import test_data, train_data, val_data
from gragh_tools import draw
from torch_geometric.data.data import Data
from create_alarms_data_002 import ORIGIN_DATA
from gragh_tools import draw

device = "cuda:0" if torch.cuda.is_available() else "cpu"


def get_sentence_vectors(texts):
    word_embeddings_list = []
    for text in texts:
        encoded_input = FTModel.embedding_text_to_tensor(text)
        word_embeddings_list.append(encoded_input.unsqueeze(0))
    return torch.cat(word_embeddings_list, dim=0)


def translate_batch(batch):
    _btree_list = [ORIGIN_DATA.btrees_origin_infos[id] for id in batch["btree"].node_id.tolist()]
    # _host_list = [ORIGIN_DATA.ips_origin_infos[id] for id in batch["host"].node_id.tolist()]
    _alarm_list = [ORIGIN_DATA.alarms_origin_infos[id] for id in batch["alarm"].node_id.tolist()]
    # print(_btree_list)
    # print(_host_list)
    # print(_alarm_list)
    batch["btree"].x = get_sentence_vectors(_btree_list)
    # batch["host"].x = get_sentence_vectors(_host_list)
    batch["alarm"].x = get_sentence_vectors(_alarm_list)
    _batch = batch.to(device)
    return _batch


MODEL_PATH = "/home/Dyf/code/storage_models/alarms/"
model = torch.load(os.path.join(MODEL_PATH, 'Alarms_{}.pth').format(28))
model = model.to(device)
model.eval()


@torch.no_grad()
def test(data):
    model.eval()
    pred = model(data)
    pred = pred.clamp(min=0, max=1)
    target = data[("alarm", "to", "btree")].edge_label.float()
    print("###############################")
    print(pred)
    print(target)
    print("###############################")
    rmse = F.mse_loss(pred, target).sqrt()
    return float(rmse)


import torch.nn.functional as F


def translate_data(data):
    new_data = translate_batch(data)
    return new_data


train_data = translate_data(train_data)
val_data = translate_data(val_data)
test_data = translate_data(test_data)

train_rmse = test(train_data)
val_rmse = test(val_data)
test_rmse = test(test_data)
print(train_rmse)
print(val_rmse)
print(test_rmse)

